import torch
import torch.cuda
import torch.nn
import torchvision.models as models
import torchvision.transforms as transforms
import numpy as np
import cv2


def get_model():
    model_path = "./models/dataset_dir_model_529.pt"
    model = torch.load(model_path )
    return model

def deal_img(img_path):
    """Transforming images on GPU"""
    image = cv2.imread(img_path) 
    image_new =  cv2.resize(image, (224,224))
    my_transforms= transforms.Compose(
        [ 
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229,0.224,0.225]) 
        ]
        )
    my_tensor = my_transforms(image_new)
    my_tensor = my_tensor.resize_(1,3,224,224)
    my_tensor= my_tensor.cuda()
    return my_tensor

def cls_inference(cls_model,imgpth):
    input_tensor = deal_img(imgpth)
    cls_model.eval()
    result = cls_model(input_tensor)
    result_npy = result.data.cpu().numpy()
    max_index = np.argmax(result_npy[0])
    return max_index

def feature_extract(cls_model,imgpth):
    cls_model.fc = torch.nn.LeakyReLU(0.1)
    cls_model.eval()
    input_tensor = deal_img(imgpth)
    result = cls_model(input_tensor)
    result_npy = result.data.cpu().numpy()
    return result_npy[0]


if __name__ == "__main__":
    image_path="./dataset_dir/train/whole_dress/dress_abstract (7).jpg"
    model = get_model()
    cls_label = cls_inference(model,image_path)
    print(cls_label)
    # feature = feature_extract(model,image_path)
    # print(feature)